Skip to content

Feat(psi): add PsiStorageMode and device_memory_mode for GPU with PW base - #7746

Open
dyzheng wants to merge 7 commits into
deepmodeling:developfrom
dyzheng:feat/issue-7727-gpu-paged
Open

Feat(psi): add PsiStorageMode and device_memory_mode for GPU with PW base#7746
dyzheng wants to merge 7 commits into
deepmodeling:developfrom
dyzheng:feat/issue-7727-gpu-paged

Conversation

@dyzheng

@dyzheng dyzheng commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Reminder

  • I have read AGENTS.md and docs/developers_guide/agent_governance.md.
  • I have linked an issue or explained why this PR does not need one.
  • I have added adequate unit tests and/or case tests, or explained why not.
  • I have listed the exact verification commands run and their results.
  • I have described user-visible behavior changes, including INPUT parameter changes.
  • I have explained core-module impact for ESolver, HSolver, ElecState, Hamilt, Operator, Psi, or other source/ changes.
  • I have requested any needed governance exception below.

Linked Issue

Fix #7727

Unit Tests and/or Case Tests for my changes

  • Commands run:
  • Result summary:
  • Checks not run, with reason:

What's changed?

  • Example: brief summary of the user-visible or developer-facing change.

Governance Notes

  • INPUT/docs changes:
  • Core module impact:
  • Exceptions requested:

dyzheng added 2 commits August 1, 2026 21:43
…paging

- Add PsiStorageMode enum (ALL_GPU, ALL_CPU, PAGED_GPU) to psi namespace
- Add paging interface: load_k_to_gpu, store_k_from_gpu, ensure_k_on_gpu
- Modify resize() to allocate CPU storage + single-k GPU buffer in PAGED_GPU mode
- Create psi_paging.cpp with paging method implementations
- Add device_memory_mode INPUT parameter (''/full_gpu/paged)
- Add device_memory_mode reading in read_input_item_system.cpp
…aths

- hsolver_pw.cpp: load_k_to_gpu/store_k_from_gpu around k-point solve loop
  Skip k-parent propagation in PAGED_GPU mode
- psi_prepare.cpp: PAGED_GPU memory handling in initialize_psi
- setup_psi_pw.cpp: detect paging mode, log GPU memory mode, copy_d2h paging path
- forces/stress onsite/nl/kin: ensure_k_on_gpu before k-point loops
- onsite_proj.cpp: ensure_k_on_gpu in cal_occupations
@mohanchen mohanchen added Bugs Bugs that only solvable with sufficient knowledge of DFT GPU & DCU & HPC GPU and DCU and HPC related any issues labels Aug 2, 2026
dyzheng added 5 commits August 2, 2026 15:15
- Fix link errors: add psi_paging.cpp to 3 test CMakeLists.txt targets
  (MODULE_ESTATE_elecstate_base, MODULE_ESTATE_elecstate_pw,
   MODULE_IO_read_wfc_nao_test, MODULE_PSI_Unittests)
- Fix SIGSEGV: use custom read_value for device_memory_mode to handle
  empty string (avoid str_values[0] out-of-bounds on empty value re-read)
- Fix stack smashing: remove unused __CUDA/__ROCM guarded members
  (compute_stream_, transfer_stream_) from Psi class to eliminate
  ABI inconsistency between library and test TUs
…ain.md

CI check-documentation step requires docs/parameters.yaml to match the
output of --generate-parameters-yaml. The PR added device_memory_mode
Input_Item but docs were not updated.
When device_memory_mode defaults to 'paged' (auto mode on GPU PW),
fix_k() returned CPU buffer pointer for k-points not yet loaded,
causing psiToRho and other GPU consumers to read CPU memory instead
of GPU memory. This produced wrong charge density and incorrect total
energy for multi-k-point GPU tests (003, 078, etc.).

Fix: make current_k_gpu_ mutable, call load_k_to_gpu in fix_k/fix_b/
fix_kb when the requested k-point differs from the loaded one. Non-PAGED
modes are unaffected (load_k_to_gpu is a no-op for ALL_GPU).
…safe

Pyabacus diagopack compiles psi.cpp directly and needs psi_paging.cpp
for the get_cpu_pointer_safe / set_storage_mode / load_k_to_gpu symbols.
install(DIRECTORY ...) only copies at install time, not during build.
The test runs from CMAKE_CURRENT_BINARY_DIR and expects ./support/
with orbital/pseudo files. Use file(COPY ...) to make them available
during ctest.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Bugs Bugs that only solvable with sufficient knowledge of DFT GPU & DCU & HPC GPU and DCU and HPC related any issues

Projects

None yet

Development

Successfully merging this pull request may close these issues.

GPU fixes and optimization for Diago_DavSubspace, plus paged GPU memory mode (PAGED_GPU)

2 participants